import os
import json
import numpy as np
from tqdm import tqdm

def parse_json(path):
    with open(path, "r") as f:
        data = json.load(f)

    T = len(data[0])
    n_agents = len(data)

    MAX_T = min(400, T)

    kills = np.zeros((MAX_T, n_agents))
    is_alive = np.ones((MAX_T, n_agents))

    for agent_id, traj in enumerate(data):
        for t, step in enumerate(traj[:MAX_T]):
            try:
                kill_count = step[6][0]
                alive_flag = step[7][0]
                kills[t, agent_id] = kill_count
                is_alive[t, agent_id] = alive_flag
            except Exception as e:
                print(f"[ERROR] {os.path.basename(path)} step {t} agent {agent_id}: {e}")
                return None
    return kills, is_alive

def check_kd(path):
    result = {
        "file": os.path.basename(path),
        "kills_found": False,
        "deaths_found": False,
        "total_kills": 0,
        "total_deaths": 0
    }

    parsed = parse_json(path)
    if parsed is None:
        return result  # 格式错误或异常跳过

    kills, is_alive = parsed
    final_kills = kills[-1]
    death_mask = ((is_alive[:-1] == 1) & (is_alive[1:] == 0)).sum()

    result["kills_found"] = np.any(final_kills > 0)
    result["deaths_found"] = death_mask > 0
    result["total_kills"] = final_kills.sum()
    result["total_deaths"] = death_mask

    return result

def main(json_dir):
    files = [f for f in os.listdir(json_dir) if f.endswith(".json")]
    all_results = []

    print(f"🔍 正在检查 {len(files)} 个 JSON 文件...\n")
    for fname in tqdm(files, desc="Checking"):
        path = os.path.join(json_dir, fname)
        res = check_kd(path)
        all_results.append(res)

    # 输出摘要
    no_kill = [r["file"] for r in all_results if not r["kills_found"]]
    no_death = [r["file"] for r in all_results if not r["deaths_found"]]

    print("\n✅ 检查完成！")
    print(f"共有 {len(files)} 个文件")
    print(f"👉 有击杀记录的文件数: {len(files) - len(no_kill)}")
    print(f"👉 有死亡记录的文件数: {len(files) - len(no_death)}")

    if no_kill:
        print("\n⚠️ 没有击杀记录的文件：")
        for f in no_kill:
            print("  -", f)

    if no_death:
        print("\n⚠️ 没有死亡记录的文件：")
        for f in no_death:
            print("  -", f)

    total_kills = sum([r["total_kills"] for r in all_results])
    total_deaths = sum([r["total_deaths"] for r in all_results])
    print(f"\n📊 总击杀数: {total_kills:.0f}, 总死亡数: {total_deaths:.0f}")

if __name__ == "__main__":
    # 修改为你的 JSON 路径
    json_folder = "offline_data/2s3z/good"
    main(json_folder)
